# matchmove.py

import numpy as np
import transform
import math
from sys import float_info

DOF = 6 		# the number of degrees of freedom for our camera

# CCD / line sweep constants
# TODO: tune all of these 
INIT_STEP_SIZES = [
	math.pi/16.0,
	math.pi/16.0,
	math.pi/16.0,
	1.0,
	1.0,
	1.0]

# ZRot, YRot, XRot, ZTrans, YTrans, XTrans
INIT_PARAM_GUESS = [
	0.0,
	0.0, 
	0.0,
	-150.0,
	-32.0,
	-39.0]

EPSILON = 1.3
STEP_SCALE = 2.0
SEARCH_LIMIT = 20
SWEEP_LIMIT = 1000

# transformMatrixDict is used to look up the function for generating the matrix
# for the given param index 
transformFunctions = [ 
	transform.rotate_z,
	transform.rotate_y,
	transform.rotate_x,
	transform.translate_z,
	transform.translate_y,	
	transform.translate_x]

# solve the matchmove problem of recovering the transforms used to position
# the camera in order to get the NDC marker positions p given the real world known
# marker positions.
# perspProjM: perspective projection matrix specific to camera used to take image
# pValList: list of NDC (x,y) for marker positions in the image
# xValList: list of 3D coordinates (x,y,z,1) for marker positions in the scene
# returns the transform matrix for the camera
def solve(perspProjM, pList, xList):
	"""
	solve performs CCD on the parameters for each component transform matrix
	and returns the optimal transformation matrix that approximates:
	p_i = transformM * x_i for each each corresponding p_i, x_i pair 	
	"""
	assert len(pList) == len(xList), "pList and xList must be of the same length"
	
	initTransformParamList = list(INIT_PARAM_GUESS)
			
	optimalTransformParamList = cyclical_coordinate_descent(perspProjM, initTransformParamList, pList, xList)
	return generate_transform(optimalTransformParamList)
# 
def find_residual(pList, newPList):
	"""
	find_residual compares the values in pList and newPList and returns the
	squared magnitude for least squares minimization
	"""
	diffList = list()
	for i in range(len(pList)):
		curDiff = pList[i] - newPList[i]
		curDiff = np.absolute(curDiff)		
		diffList.append(np.sum(curDiff))	
	
	diffArray = np.array(diffList)
	residual = np.sqrt(np.vdot(diffArray, diffArray))
	return residual	

def generate_transform(transformParamList):
	"""
	generate_transform creates a transform matrix which the result of 
	matrix multiples of the matrices created by calling each function
	in transformMatrixDict with the corresponding param 	
	"""
	transformList = list()
	for i in range(len(transformParamList)):
		param = transformParamList[i]
		curT = transformFunctions[i](param)
		transformList.append(curT)
	
	return transform.composeTransformList(transformList)	

def eval_transform(perspProjM, transformParamList, xList):
	"""
	eval_transform does perspProjM * transformM * xValArray[i] for each set of 3D marker scene positions (xVals)
	returns a new array of pVals (for each 2D marker positions in image)	
	"""
	# build the transform matrix resulting from the sequence of matrix multiplies: R_z * R_y * R_x * T_z * T_y * T_x
	mvTransform = generate_transform(transformParamList)
			
	# multiply the transformation matrix by the perspective projection matrix to get the final transformation matrix
	pmvTransform = transform.compose(perspProjM, mvTransform)
	
	# multiply each 3D scene marker position by the final transformation matrix to get the new NDC marker position
	newPList = list()
	for i in range(len(xList)):		
		newP = np.dot(pmvTransform, xList[i])
		newPNDC = newP / newP[3]
		newPList.append(np.array([newPNDC[0], newPNDC[1]]))
		
	return newPList

def eval_objective(perspProjM, transformParamList, pList, xList):
	"""
	finds the residual for the given param list
	"""
	# evaluate the transform for each 3D scene marker position: perspProjM * transformM * x_i 
	# to get a new list of p values for comparing	
	newPList = eval_transform(perspProjM, transformParamList, xList)
		
	# now compare the known list of p values to the p values generated from applying
	# the test transform 
	residual = find_residual(pList, newPList)
	
	return residual	

def find_search_direction(i, perspProjM, transformParamList, pList, xList):
	# try positive step
	val = transformParamList[i]
	transformParamList[i] = val + INIT_STEP_SIZES[i]
	posResidual = eval_objective(perspProjM, transformParamList, pList, xList)
	transformParamList[i] = val - INIT_STEP_SIZES[i]
	negResidual = eval_objective(perspProjM, transformParamList, pList, xList)
	
	searchDirection = 1.0	
	if negResidual < posResidual:
		searchDirection = searchDirection * -1.0		
		
	return searchDirection	

def find_step_scale(i, perspProjM, transformParamList, pList, xList, searchDirection, curStepSize):
	val = transformParamList[i]
	transformParamList[i] = val + 2.0 * searchDirection * curStepSize
	doubleResidual = eval_objective(perspProjM, transformParamList, pList, xList)
	transformParamList[i] = val + 0.5 * searchDirection * curStepSize
	halfResidual = eval_objective(perspProjM, transformParamList, pList, xList)
	
	scale = 2.0
	if halfResidual < doubleResidual:
		scale = 0.5	
		
	return scale
	
def line_search(i, perspProjM, transformParamList, pList, xList, initResidual, searchDirection):
	"""
	line search iterations on the param in transformParamList given by the
	param index i
	this implementation uses step size scaling and back tracking in order to
	descend on the best value for param i, given the current values of the other 
	parameters in a constrained number of searches	
	"""	
	
	numSearches = 0
	stepSize = INIT_STEP_SIZES[i]
	initVal = transformParamList[i]
	curResidual = initResidual
	while numSearches < SEARCH_LIMIT and curResidual >= initResidual:
		# save the cur value of param i
		curVal = initVal
	
		# need to determine search direction before continuing
		oneStepVal = curVal + searchDirection * stepSize
		twoStepVal = curVal + searchDirection * 2.0 * stepSize
		oneAndHalfStepVal = curVal + -1.0 * searchDirection * stepSize		
		
		# find out which step is the best to take
		transformParamList[i] = oneStepVal		
		oneStepResidual = eval_objective(perspProjM, list(transformParamList), pList, xList)
		
		transformParamList[i] = twoStepVal
		twoStepResidual = eval_objective(perspProjM, list(transformParamList), pList, xList)
		
		transformParamList[i] = oneAndHalfStepVal
		oneAndHalfStepResidual = eval_objective(perspProjM, list(transformParamList), pList, xList)
		
		curResidual = oneStepResidual
		curVal = oneStepVal
		if twoStepResidual < curResidual:
			curResidual = twoStepResidual
			curVal = twoStepVal
			
		if oneAndHalfStepResidual < curResidual:
			curResidual = oneAndHalfStepResidual
			curVal = oneAndHalfStepVal				
		
		numSearches = numSearches + 1

	val = initVal
	residual = initResidual			
	if curResidual < initResidual:
		val = curVal
		residual = curResidual
		
	return val, residual
		
def cyclical_coordinate_descent(perspProjM, transformParamList, pList, xList):
	"""
	performs sweeps over the parameters for the various component transform matrices,
	using line search to optimize each parameter in a given sweep
	"""
	
	# determine the search directions for each parameter
	searchDirections = list()
	for i in range(len(transformParamList)):
		curDirection = find_search_direction(i, perspProjM, list(transformParamList), pList, xList)
		searchDirections.append(curDirection)
		# print "direction for param i: %d" % (curDirection)
	
	prevResidual = eval_objective(perspProjM, transformParamList, pList, xList)
	curResidual = prevResidual
	prevTransformParamList = transformParamList
	curTransformParamList = list(prevTransformParamList)		
	# CCD sweeps
	num_sweeps = 0
	while prevResidual > EPSILON:
		prevTransformParamList = list(curTransformParamList)		
		# print prevTransformParamList
		for i in range(3, len(transformParamList)):
			# do a line search on parameter i
			(curVal, curResidual) = line_search(i, perspProjM, list(transformParamList), pList, xList, curResidual, searchDirections[i])			
			transformParamList[i] = curVal
			prevResidual = curResidual
			
		# now do a linear search along the vector difference between prev and current
		testParamList = list(prevTransformParamList)
		for i in range(1, 10):
			for i in range(len(transformParamList)):
				testParamList[i] = prevTransformParamList[i] + i*0.10*(curTransformParamList[i] - prevTransformParamList[i])
			
			testResidual = eval_objective(perspProjM, testParamList, pList, xList)
			if testResidual < prevResidual:				
				curTransformParamList = list(testParamList)
				prevResidual = testResidual
				break
			
		num_sweeps = num_sweeps + 1
		
	return prevTransformParamList

'''
def test():
	# create a list of arrays for the x values
	xList = list()
	xList.append(np.array([0.,  0.,   0., 1]))
	xList.append(np.array([0.,  165., 0., 1]))
	xList.append(np.array([80., 165., 0., 1]))
	xList.append(np.array([80., 0.,   0., 1.]))

	# NDC for each marker
	pList = list() 
	pList.append(np.array([-0.925, -0.75555556]))
	pList.append(np.array([-0.284375, 0.05]))
	pList.append(np.array([0.353125, 0.01666667]))
	pList.append(np.array([0.934375, -0.8]))
		
	# setup the perspective projection matrix
	fovy = 0.8178787407770628
	aspect = 1.504
	z_near = 45.0
	z_far = 200.0
	perspProjM = transform.perspective(fovy, aspect, z_near, z_far)
	transformM = solve(perspProjM, pList, xList)
	
	print transformM 
	
	newPList = transform.compose(perspProjM, transformM, xList)
	for i in range(len(newPList)):
		newPList[i] = newPList[i]/newPList[i][3]	
	
	print newPList
	print PList
# Execute
if __name__ == "__main__":
	test()
	
'''
